/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*!
 * \file flash_attention_score_grad_tiling_g_bn2.cc
 * \brief
 */

#include "flash_attention_score_grad_tiling_g_bn2.h"
#include "tiling/tiling_type.h"
#include "tiling/tiling_templates_registry.h"

using namespace AscendC;
using namespace ge;

namespace optiling {

constexpr uint32_t BYTE_BLOCK = 32;
constexpr uint32_t FRACTAL_NUM = 16; // 分形大小

constexpr uint32_t gScaleAttrIdx     = 0;
constexpr uint32_t gKeepProbAttrIdx  = 1;
constexpr uint32_t gPrevTokenAttrIdx = 2;
constexpr uint32_t gNextTokenAttrIdx = 3;
constexpr uint32_t gHeadNumAttrIdx   = 4;
constexpr uint32_t gLayoutAttrIdx    = 5;

#define KB        1024
#define UB_SIZE   (192 * KB)
#define L1_SIZE   (512 * KB)
#define L0AB_SIZE (64 * KB)
#define L0C_SIZE  (128 * KB)

bool FlashAttentionScoreGradUgBbn::IsCapable() {
  return false; //zyt temp
}
uint64_t FlashAttentionScoreGradUgBbn::GetTilingKey() const {
  auto queryType = context_->GetInputTensor(0)->GetDataType();
  if (queryType == ge::DT_FLOAT16) {
    // 默认 fp16 使用高精度模式
    return TILINGKEY(NONE, G, N2, FLOAT16, BSND, ALL);
  } else if (queryType == ge::DT_BF16) {
    // bfloat16
    return TILINGKEY(NONE, G, N2, BFLOAT16, BSND, ALL);
  } else if (queryType == ge::DT_FLOAT) {
    // float32
    return TILINGKEY(NONE, G, N2, FLOAT32, BSND, ALL);
  } else {
    return TILINGKEY(NONE, G, N2, FLOAT16, BSND, ALL);
  }
}

uint32_t CeilFuncV1(uint32_t num1, uint32_t num2)
{
    if (num2 == 0) {
        return 0;
    }
    return (num1 + num2 - 1) / num2;
}

ge::graphStatus FlashAttentionScoreGradUgBbn::GetPlatformInfo() {
  // 1. Get CoreNum
  auto platformInfoPtr = context_->GetPlatformInfo();
  if (platformInfoPtr == nullptr) {
    auto compileInfoPtr = reinterpret_cast<const FlashAttentionScoreGradCompileInfo*>(context_->GetCompileInfo());
    OPS_ERR_IF(compileInfoPtr == nullptr, OPS_REPORT_CUBE_INNER_ERR(context_, "compile_info is null"),
               return ge::GRAPH_FAILED);

    mBaseInfoParams.coreNum = compileInfoPtr->aivNum;
    mBaseInfoParams.aicNum = compileInfoPtr->aicNum;
  } else {
    ascendcPlatform_ = std::unique_ptr<platform_ascendc::PlatformAscendC>(
        new (std::nothrow) platform_ascendc::PlatformAscendC(platformInfoPtr));
    mBaseInfoParams.coreNum = ascendcPlatform_->GetCoreNumAiv();
    mBaseInfoParams.aicNum = ascendcPlatform_->GetCoreNumAic();
  }

  // 2. Get buffer size
  mBaseInfoParams.ubSize = UB_SIZE;
  mBaseInfoParams.l1Size = L1_SIZE;
  mBaseInfoParams.l0aSize = L0AB_SIZE;
  mBaseInfoParams.l0bSize = L0AB_SIZE;
  mBaseInfoParams.l0cSize = L0C_SIZE;

  return ge::GRAPH_SUCCESS;
}
ge::graphStatus FlashAttentionScoreGradUgBbn::GetShapeAttrsInfo() {
  uint32_t b = 0;
  uint32_t n1 = 0;
  uint32_t n2 = 0;
  uint32_t d = 0;
  uint32_t h = 0;
  uint32_t s1 = 0;
  uint32_t s2 = 0;
  uint32_t g = 1;
  uint32_t layoutType = 0;

  const char* inputLayout = context_->GetAttrs()->GetAttrPointer<char>(gLayoutAttrIdx);
  const gert::StorageShape* queryShape = context_->GetInputShape(0);
  const gert::StorageShape* keyShape = context_->GetInputShape(1);

  b  = queryShape->GetStorageShape().GetDim(0);
  s1 = queryShape->GetStorageShape().GetDim(1);
  h  = queryShape->GetStorageShape().GetDim(2);
  n1 = *context_->GetAttrs()->GetAttrPointer<int>(gHeadNumAttrIdx);
  n2 = n1;
  if (inputLayout[0] == 'B' && inputLayout[1] == 'S' && inputLayout[2] == 'H') {
      s2 = keyShape->GetStorageShape().GetDim(1);
      d = h / n1;
  } else if (inputLayout[0] == 'S' && inputLayout[1] == 'B' && inputLayout[2] == 'H') {
      b  = queryShape->GetStorageShape().GetDim(1);
      s1 = queryShape->GetStorageShape().GetDim(0);
      s2 = keyShape->GetStorageShape().GetDim(0);
      d = h / n1;
      layoutType = 1;
  } else if (inputLayout[0] == 'B' && inputLayout[1] == 'N' && inputLayout[2] == 'S' && inputLayout[3] == 'D') {
      n1 = queryShape->GetStorageShape().GetDim(1);
      s1 = queryShape->GetStorageShape().GetDim(2);
      n2 = keyShape->GetStorageShape().GetDim(1);
      s2 = keyShape->GetStorageShape().GetDim(2);
      d = queryShape->GetStorageShape().GetDim(3);
      g  = n1 / n2;
      layoutType = 2;
  } else if (inputLayout[0] == 'B' && inputLayout[1] == 'S' && inputLayout[2] == 'N' && inputLayout[3] == 'D') {
      n1 = queryShape->GetStorageShape().GetDim(2);
      s2 = keyShape->GetStorageShape().GetDim(1);
      n2 = keyShape->GetStorageShape().GetDim(2);
      d = queryShape->GetStorageShape().GetDim(3);
      g = n1 / n2;
      layoutType = 3;
  }

  mBaseInfoParams.b      = b;
  mBaseInfoParams.n1     = n1;
  mBaseInfoParams.n2     = n2;
  mBaseInfoParams.s1     = s1;
  mBaseInfoParams.s2     = s2;
  mBaseInfoParams.d      = d;
  mBaseInfoParams.g      = g;
  mBaseInfoParams.h      = h;
  mBaseInfoParams.layout = layoutType;
  // sparse
  int preTokens  = *context_->GetAttrs()->GetAttrPointer<int>(gPrevTokenAttrIdx);
  int nextTokens = *context_->GetAttrs()->GetAttrPointer<int>(gNextTokenAttrIdx);
  mBaseInfoParams.isSparse = ((uint32_t)preTokens < mBaseInfoParams.s1 || (uint32_t)nextTokens < mBaseInfoParams.s1) ? true : false;

  // 数据类型
  auto queryType = context_->GetInputTensor(0)->GetDataType();
  if (queryType == ge::DT_FLOAT16) {
      mBaseInfoParams.dataTypeSize = 2;
      mBaseInfoParams.mask         = 128;
      mBaseInfoParams.dataType     = 1;
  } else if (queryType == ge::DT_BF16) {
      mBaseInfoParams.dataTypeSize = 2;
      mBaseInfoParams.mask         = 64;
      mBaseInfoParams.dataType     = 2;
  } else {
      mBaseInfoParams.dataTypeSize = 4;
      mBaseInfoParams.mask         = 64;
      mBaseInfoParams.dataType     = 0;
  }

  mBaseInfoParams.calTypeSize = 4;
  mBaseInfoParams.dType = queryType;
  mBaseInfoParams.blockNum = BYTE_BLOCK / mBaseInfoParams.dataTypeSize;
  return ge::GRAPH_SUCCESS;
}

ge::graphStatus FlashAttentionScoreGradUgBbn::SetBaseTilingData() {
  // base tiling data
  tilingData.gBNBaseParams.set_coreNum(mBaseInfoParams.coreNum);

  tilingData.gBNBaseParams.set_B(mBaseInfoParams.b);
  tilingData.gBNBaseParams.set_N1(mBaseInfoParams.n1);
  tilingData.gBNBaseParams.set_N2(mBaseInfoParams.n2);
  tilingData.gBNBaseParams.set_S1(mBaseInfoParams.s1);
  tilingData.gBNBaseParams.set_S2(mBaseInfoParams.s2);
  tilingData.gBNBaseParams.set_D(mBaseInfoParams.d);
  tilingData.gBNBaseParams.set_G(mBaseInfoParams.g);
  tilingData.gBNBaseParams.set_H(mBaseInfoParams.h);
  tilingData.gBNBaseParams.set_batch(mBaseInfoParams.b * mBaseInfoParams.n1);  // bn开多核场景下，batch为BN
  tilingData.gBNBaseParams.set_layout(mBaseInfoParams.layout);

  tilingData.gBNBaseParams.set_typeBlockNum(32 / sizeof(float));
  tilingData.gBNBaseParams.set_scaleValue(*(context_->GetAttrs()->GetAttrPointer<float>(gScaleAttrIdx)));
  tilingData.gBNBaseParams.set_keepProb(*(context_->GetAttrs()->GetAttrPointer<float>(gKeepProbAttrIdx)));

  int preTokens  = *context_->GetAttrs()->GetAttrPointer<int>(gPrevTokenAttrIdx);
  int nextTokens = *context_->GetAttrs()->GetAttrPointer<int>(gNextTokenAttrIdx);
  tilingData.gBNBaseParams.set_preTokens(preTokens);
  tilingData.gBNBaseParams.set_nextTokens(nextTokens);

  return ge::GRAPH_SUCCESS;
}

ge::graphStatus FlashAttentionScoreGradUgBbn::SetCoresSplitTilingData() {
  uint32_t gInner = mBaseInfoParams.gInner;
  uint32_t gOut = CeilFuncV1(mBaseInfoParams.g, mBaseInfoParams.gInner);
  uint32_t gTail = mBaseInfoParams.g % mBaseInfoParams.gInner;
  uint32_t innerMatResSize = mBaseInfoParams.s1 * mBaseInfoParams.s2 * mBaseInfoParams.calTypeSize;

  tilingData.gBNSplitCoreParams.set_gOut(gOut);
  tilingData.gBNSplitCoreParams.set_gInner(gInner);
  tilingData.gBNSplitCoreParams.set_gTail(gTail);
  tilingData.gBNSplitCoreParams.set_innerMatResSize(innerMatResSize);
  return ge::GRAPH_SUCCESS;
}

ge::graphStatus FlashAttentionScoreGradUgBbn::SetUbSplitTilingData() {
    uint32_t batch = tilingData.gBNBaseParams.get_batch();

    uint32_t batchRange = CeilFuncV1(batch, mBaseInfoParams.coreNum); // ceil
    uint32_t usedCoreNum = CeilFuncV1(batch, batchRange);
    uint32_t remainBatchRange = batch - batchRange * (usedCoreNum - 1);

    mBaseInfoParams.coreNum = usedCoreNum;
    tilingData.gBNBaseParams.set_coreNum(mBaseInfoParams.coreNum);

    tilingData.gBNSplitCoreParams.set_batchRange(batchRange);        // 开核size
    tilingData.gBNSplitCoreParams.set_splitedBatchRange(batchRange); // 每个核处理总大小
    tilingData.gBNSplitCoreParams.set_splitedRemainBatchRange(remainBatchRange);
    return ge::GRAPH_SUCCESS;
}

// gInner 从 256 逐步缩小, 找到适合UB大小的最大 gInner
uint32_t FlashAttentionScoreGradUgBbn::CalcGInner()
{
    uint32_t gInner = std::min(mBaseInfoParams.g, uint32_t(256 -16));
    while (!CheckGInnerLegal(gInner)) {
        if (gInner <= FRACTAL_NUM) {
            return 0;
        }
        gInner = gInner - FRACTAL_NUM;
    }
    return gInner;
};

bool FlashAttentionScoreGradUgBbn::CheckGInnerLegal(uint32_t gInner)
{
    uint32_t coexNode = 6; // 6: 3 TQue, 3 Tbuff

    // input scene
    uint32_t baseS1S2 = gInner * mBaseInfoParams.s1 * mBaseInfoParams.s2 * mBaseInfoParams.calTypeSize;
    uint32_t baseS1D  = gInner * mBaseInfoParams.s1 * mBaseInfoParams.d * mBaseInfoParams.calTypeSize;
    uint32_t baseS2D  = gInner * mBaseInfoParams.s2 * mBaseInfoParams.d * mBaseInfoParams.calTypeSize;
    uint32_t maxBufferSize = std::max(baseS1S2, std::max(baseS1D, baseS2D));

    // simplesoftmax and dropout
    auto softmaxShape = Shape({mBaseInfoParams.s1, mBaseInfoParams.s2});
    auto dropoutShape = Shape({mBaseInfoParams.s1, mBaseInfoParams.s2});
    auto softmaxgradShape = Shape({mBaseInfoParams.s1, mBaseInfoParams.s2});
    uint32_t softmaxTmpSize = GetSoftMaxMinTmpSize(softmaxShape, mBaseInfoParams.calTypeSize, true);
    uint32_t dropoutTmpSize = GetDropOutMinTmpSize(dropoutShape, mBaseInfoParams.calTypeSize, true);
    uint32_t softmaxgradTmpSize = GetSoftMaxGradMinTmpSize(softmaxgradShape, mBaseInfoParams.calTypeSize, true, false);
    uint32_t maxTmpBufferSize = std::max(softmaxTmpSize, std::max(dropoutTmpSize, softmaxgradTmpSize));

    maxBufferSize = std::max(maxTmpBufferSize, maxBufferSize);
    uint32_t baseUB = coexNode * maxBufferSize;
    if (baseUB > mBaseInfoParams.ubSize) {
        return false;
    }

    // Loc buffer
    if (baseS1S2 > mBaseInfoParams.l0cSize) {
        return false;
    }

    return true;
}

ge::graphStatus FlashAttentionScoreGradUgBbn::DoOpTiling() {
  SetBaseTilingData();
  SetCoresSplitTilingData();
  SetUbSplitTilingData();
  return ge::GRAPH_SUCCESS;
}

ge::graphStatus FlashAttentionScoreGradUgBbn::DoLibApiTiling() {
  // mm tiling
  matmul_tiling::MatmulApiTiling mm1;
  matmul_tiling::MatmulApiTiling mm2;
  matmul_tiling::MatmulApiTiling mm3;

  // s轴不切，所以默认singleM与originM等大，一个入参
  // BMM1 BMM2
  SetMatmulTiling1(mm1, tilingData.mm1TilingData);

  // BMM4 BMM3.2
  SetMatmulTiling2(mm2, tilingData.mm2TilingData);

  // BMM3.1
  SetMatmulTiling3(mm3, tilingData.mm3TilingData);

  // softmax
  auto softmaxShape = Shape({mBaseInfoParams.s1, mBaseInfoParams.s2});
  SoftMaxTilingFunc(softmaxShape, sizeof(float), (mBaseInfoParams.s1 * mBaseInfoParams.s2),
                    tilingData.softmaxTilingData);

  // softmaxgrad
  SoftMaxGradTilingFunc(softmaxShape, mBaseInfoParams.calTypeSize,
                        (mBaseInfoParams.s1 * mBaseInfoParams.s2),
                        tilingData.softmaxGradTilingData, true);
  return ge::GRAPH_SUCCESS;
}

void FlashAttentionScoreGradUgBbn::SetMatmulTiling1(matmul_tiling::MatmulApiTiling& mm, TCubeTiling& mmTiling)
{
    // set matmul type
    mm.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND,
                matmul_tiling::DataType::DT_BFLOAT16, false);
    mm.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND,
                matmul_tiling::DataType::DT_BFLOAT16, true);
    mm.SetCType(matmul_tiling::TPosition::VECCALC, matmul_tiling::CubeFormat::ND,
                matmul_tiling::DataType::DT_FLOAT);

    // set base shape
    mm.SetShape(mBaseInfoParams.s1, mBaseInfoParams.s2, mBaseInfoParams.d);
    uint32_t alignNBlockNum = (mBaseInfoParams.s2 + mBaseInfoParams.blockNum -1) / mBaseInfoParams.blockNum;
    uint32_t alignNSize = alignNBlockNum * mBaseInfoParams.blockNum;
    // set origin shape
    if (mBaseInfoParams.layout == 0 || mBaseInfoParams.layout == 3) {    // 0:BSH or 3:BSND
        mm.SetOrgShape(mBaseInfoParams.s1, alignNSize,
                       mBaseInfoParams.d * mBaseInfoParams.n2 * mBaseInfoParams.g,
                       mBaseInfoParams.d * mBaseInfoParams.n2);
    } else if (mBaseInfoParams.layout == 1) {   // 1:SBH
        mm.SetOrgShape(mBaseInfoParams.s1, alignNSize,
                       mBaseInfoParams.d * mBaseInfoParams.b * mBaseInfoParams.n2 * mBaseInfoParams.g,
                       mBaseInfoParams.d * mBaseInfoParams.b * mBaseInfoParams.n2);
    } else {  // BNSD
        mm.SetOrgShape(mBaseInfoParams.s1, alignNSize, mBaseInfoParams.d);
    }
    mm.SetBias(false);
    mm.SetBufferSpace(-1, -1, mBaseInfoParams.s1 * alignNSize);   // L1, L0C, UB, bt
    mm.SetFixSplit(mBaseInfoParams.s1, mBaseInfoParams.s1, -1);
    mm.GetTiling(mmTiling);

    mmTiling.set_shareMode(0);
    mmTiling.set_shareL1Size(mBaseInfoParams.l1Size);
    mmTiling.set_shareL0CSize(mBaseInfoParams.l0cSize);
}

void FlashAttentionScoreGradUgBbn::SetMatmulTiling2(matmul_tiling::MatmulApiTiling& mm, TCubeTiling& mmTiling)
{
    // set matmul type
    mm.SetAType(matmul_tiling::TPosition::VECCALC, matmul_tiling::CubeFormat::ND,
                matmul_tiling::DataType::DT_BFLOAT16, true);
    mm.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND,
                matmul_tiling::DataType::DT_BFLOAT16, true);
    mm.SetCType(matmul_tiling::TPosition::VECCALC, matmul_tiling::CubeFormat::ND,
                matmul_tiling::DataType::DT_FLOAT);

    // set base shape
    mm.SetShape(mBaseInfoParams.s2, mBaseInfoParams.d, mBaseInfoParams.s1);

    // set origin shape
    if (mBaseInfoParams.layout == 0 || mBaseInfoParams.layout == 3) {    // 0:BSH or 3:BSND
        mm.SetOrgShape(mBaseInfoParams.s2,
                       mBaseInfoParams.d * mBaseInfoParams.n2 * mBaseInfoParams.g,
                       mBaseInfoParams.s1,
                       mBaseInfoParams.s1);
    } else if (mBaseInfoParams.layout == 1) {   // 1:SBH
        mm.SetOrgShape(mBaseInfoParams.s2,
                       mBaseInfoParams.d * mBaseInfoParams.b * mBaseInfoParams.n2 * mBaseInfoParams.g,
                       mBaseInfoParams.s1,
                       mBaseInfoParams.s1);
    } else {  // BNSD
        mm.SetOrgShape(mBaseInfoParams.s2, mBaseInfoParams.d, mBaseInfoParams.s1, mBaseInfoParams.s1);
    }
    mm.SetBias(false);
    mm.SetBufferSpace(-1, -1, mBaseInfoParams.s1 * mBaseInfoParams.s1);   // L1, L0C, UB, bt
    mm.SetFixSplit(mBaseInfoParams.s2, mBaseInfoParams.d, -1);
    mm.GetTiling(mmTiling);

    mmTiling.set_shareMode(0);
    mmTiling.set_shareL1Size(mBaseInfoParams.l1Size);
    mmTiling.set_shareL0CSize(mBaseInfoParams.l0cSize);
}

void FlashAttentionScoreGradUgBbn::SetMatmulTiling3(matmul_tiling::MatmulApiTiling& mm, TCubeTiling& mmTiling)
{
    // set matmul type
    mm.SetAType(matmul_tiling::TPosition::VECCALC, matmul_tiling::CubeFormat::ND,
                matmul_tiling::DataType::DT_BFLOAT16, false);
    mm.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND,
                matmul_tiling::DataType::DT_BFLOAT16, false);
    mm.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND,
                matmul_tiling::DataType::DT_FLOAT);

    // set base shape
    mm.SetShape(mBaseInfoParams.s1, mBaseInfoParams.d, mBaseInfoParams.s2);

    // set origin shape
    if (mBaseInfoParams.layout == 0 || mBaseInfoParams.layout == 3) {    // BSH or BSND
        mm.SetOrgShape(mBaseInfoParams.s1,
                       mBaseInfoParams.d * mBaseInfoParams.n2,
                       mBaseInfoParams.s2,
                       mBaseInfoParams.s2);
    } else if (mBaseInfoParams.layout == 1) {   // SBH
        mm.SetOrgShape(mBaseInfoParams.s1,
                       mBaseInfoParams.d * mBaseInfoParams.b * mBaseInfoParams.n2,
                       mBaseInfoParams.s2,
                       mBaseInfoParams.s2);
    } else {  // BNSD
        mm.SetOrgShape(mBaseInfoParams.s1, mBaseInfoParams.d, mBaseInfoParams.s2, mBaseInfoParams.s2);
    }
    mm.SetBias(false);
    mm.SetBufferSpace(-1, -1, mBaseInfoParams.s1 * mBaseInfoParams.s2);   // L1, L0C, UB, bt
    mm.SetFixSplit(mBaseInfoParams.s1, mBaseInfoParams.d, -1);
    mm.GetTiling(mmTiling);

    mmTiling.set_shareMode(0);
    mmTiling.set_shareL1Size(mBaseInfoParams.l1Size);
    mmTiling.set_shareL0CSize(mBaseInfoParams.l0cSize);
}

ge::graphStatus FlashAttentionScoreGradUgBbn::GetWorkspaceSize() {
  auto aicNum = mBaseInfoParams.aicNum;
  auto aivNum = mBaseInfoParams.coreNum;

  uint32_t coreNum = tilingData.gBNBaseParams.get_coreNum();
  uint32_t innerMatResSize = tilingData.gBNSplitCoreParams.get_innerMatResSize();
  context_->SetBlockDim(CalcTschBlockDim(coreNum, aicNum, aivNum));
  size_t* workspaces = context_->GetWorkspaceSizes(1);
  workspaces[0] = 100 * 1024 * 1024  /* 100MB = 100 * 1024KB = 100 * 1024 * 1024 B */
                  + innerMatResSize * coreNum * coreNum * 4; // 4 is min core num

  return ge::GRAPH_SUCCESS;
}
ge::graphStatus FlashAttentionScoreGradUgBbn::PostTiling() {
  tilingData.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
  context_->GetRawTilingData()->SetDataSize(tilingData.GetDataSize());
  return ge::GRAPH_SUCCESS;
}

}  // namespace optiling
